{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-04-01T06:16:40.145999Z",
     "start_time": "2025-04-01T06:16:40.142528Z"
    }
   },
   "source": [
    "from collections import defaultdict, Counter\n",
    "from os.path import split\n",
    "\n",
    "from jinja2.lexer import compile_rules\n",
    "from sympy import lerchphi\n",
    "\n",
    "# defaultdict是一种字典，当访问不存在的键时会返回一个默认\n",
    "# Counter是一个计数器工具\n",
    "sentences = [\n",
    "    \"我\", \"喜欢\", \"吃\", \"苹果\", \"他\", \"不\", \"喜欢\", \"吃\", \"苹果派\",\n",
    "    \"I like to eat apples\", \"She has a cute cat\", \"you are very cute\", \"give you a hug\"\n",
    "]     # 句子列表"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T06:16:43.406194Z",
     "start_time": "2025-04-01T06:16:43.403356Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 定义统计函数\n",
    "def build_stats(sentences):\n",
    "    stats = defaultdict(int)    # 默认返回值为0\n",
    "    for sentence in sentences:\n",
    "        symbols = sentence.split()  # 将句子分割成单词列表(对中文的处理可能不符合预期)\n",
    "        for symbol in symbols:\n",
    "            stats[symbol] += 1\n",
    "    return stats\n",
    "\n",
    "stats = build_stats(sentences)\n",
    "print(\"stats: \", stats)"
   ],
   "id": "2ab1231b9cb516e",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "stats:  defaultdict(<class 'int'>, {'我': 1, '喜欢': 2, '吃': 2, '苹果': 1, '他': 1, '不': 1, '苹果派': 1, 'I': 1, 'like': 1, 'to': 1, 'eat': 1, 'apples': 1, 'She': 1, 'has': 1, 'a': 2, 'cute': 2, 'cat': 1, 'you': 2, 'are': 1, 'very': 1, 'give': 1, 'hug': 1})\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T06:17:00.632136Z",
     "start_time": "2025-04-01T06:17:00.627066Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 收集字母特征\n",
    "alphabet = []\n",
    "for word in stats.keys():\n",
    "    if word[0] not in alphabet:     # 收集所有单词的首字母到alphabet\n",
    "        alphabet.append(word[0])\n",
    "    for letter in word[1:]:         # 收集单词后续字符，用 ## 前缀标记-用于区分首字符与其他字符\n",
    "        if f\"##{letter}\" not in alphabet:\n",
    "            alphabet.append(f\"##{letter}\")\n",
    "alphabet.sort()\n",
    "alphabet"
   ],
   "id": "56bd70c6f618799f",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['##a',\n",
       " '##e',\n",
       " '##g',\n",
       " '##h',\n",
       " '##i',\n",
       " '##k',\n",
       " '##l',\n",
       " '##o',\n",
       " '##p',\n",
       " '##r',\n",
       " '##s',\n",
       " '##t',\n",
       " '##u',\n",
       " '##v',\n",
       " '##y',\n",
       " '##果',\n",
       " '##欢',\n",
       " '##派',\n",
       " 'I',\n",
       " 'S',\n",
       " 'a',\n",
       " 'c',\n",
       " 'e',\n",
       " 'g',\n",
       " 'h',\n",
       " 'l',\n",
       " 't',\n",
       " 'v',\n",
       " 'y',\n",
       " '不',\n",
       " '他',\n",
       " '吃',\n",
       " '喜',\n",
       " '我',\n",
       " '苹']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T06:19:29.990224Z",
     "start_time": "2025-04-01T06:19:29.986342Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 初始化词表\n",
    "vocab = alphabet.copy()\n",
    "splits = {\n",
    "    word: [c if i == 0 else f\"##{c}\" for i, c in enumerate(word)]\n",
    "    for word in stats.keys()\n",
    "}\n",
    "splits"
   ],
   "id": "17243389a65696bb",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'我': ['我'],\n",
       " '喜欢': ['喜', '##欢'],\n",
       " '吃': ['吃'],\n",
       " '苹果': ['苹', '##果'],\n",
       " '他': ['他'],\n",
       " '不': ['不'],\n",
       " '苹果派': ['苹', '##果', '##派'],\n",
       " 'I': ['I'],\n",
       " 'like': ['l', '##i', '##k', '##e'],\n",
       " 'to': ['t', '##o'],\n",
       " 'eat': ['e', '##a', '##t'],\n",
       " 'apples': ['a', '##p', '##p', '##l', '##e', '##s'],\n",
       " 'She': ['S', '##h', '##e'],\n",
       " 'has': ['h', '##a', '##s'],\n",
       " 'a': ['a'],\n",
       " 'cute': ['c', '##u', '##t', '##e'],\n",
       " 'cat': ['c', '##a', '##t'],\n",
       " 'you': ['y', '##o', '##u'],\n",
       " 'are': ['a', '##r', '##e'],\n",
       " 'very': ['v', '##e', '##r', '##y'],\n",
       " 'give': ['g', '##i', '##v', '##e'],\n",
       " 'hug': ['h', '##u', '##g']}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T06:33:45.733665Z",
     "start_time": "2025-04-01T06:33:45.729790Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def compute_pair_scores(splits):\n",
    "    \"\"\"\n",
    "    计算所有可能字符对的得分，用于WordPiece分词算法的训练阶段\n",
    "\n",
    "    参数:\n",
    "        splits: 字典{单词: 当前拆分后的字符列表}\n",
    "                例如 {\"hello\": [\"h\", \"e\", \"l\", \"l\", \"o\"], ...}\n",
    "        stats: 全局的单词频率统计字典{单词: 出现次数} (注意: 参数中未显示但实际使用)\n",
    "\n",
    "    返回:\n",
    "        字典{(char1, char2): 得分}，得分越高表示这对字符越应该合并\n",
    "    \"\"\"\n",
    "    letter_freqs = defaultdict(int) # 统计每个字符的总出现次数\n",
    "    pair_freqs = defaultdict(int)   # 统计每个相邻字符对的出现次数\n",
    "    for word, freq in stats.items():\n",
    "        split = splits[word]\n",
    "        if len(split) == 1:         # 无法再分割\n",
    "            letter_freqs[split[0]] += freq\n",
    "            continue\n",
    "        for i in range(len(split) - 1):\n",
    "            pair = (split[i], split[i + 1])\n",
    "            letter_freqs[split[i]] += freq\n",
    "            pair_freqs[pair] += freq\n",
    "        letter_freqs[split[-1]] += freq\n",
    "\n",
    "    # 计算每个字符对的得分 (核心公式)-->如果这两个letter同时出现的次数越多，score越高。\n",
    "    # score = 该字符对pair出现的频率 / 两个字符各自独立出现频率的乘积\n",
    "    scores = {\n",
    "        pair: freq / (letter_freqs[pair[0]] * letter_freqs[pair[1]])\n",
    "        for pair, freq in pair_freqs.items()\n",
    "    }   # WordPiece/BPE分词算法的核心步骤-->例:(\"h\",\"e\") 得分最高，所有单词中的\"h\"+\"e\"会被合并为\"he\"\n",
    "    return scores"
   ],
   "id": "c211667344d8cea2",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T06:39:59.076167Z",
     "start_time": "2025-04-01T06:39:59.073053Z"
    }
   },
   "cell_type": "code",
   "source": [
    "pair_scores = compute_pair_scores(splits)\n",
    "for key in list(pair_scores.keys())[:10]:\n",
    "    print(f\"{key}:{pair_scores[key]}\")"
   ],
   "id": "cad452dab8676be5",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('喜', '##欢'):0.5\n",
      "('苹', '##果'):0.5\n",
      "('##果', '##派'):0.5\n",
      "('l', '##i'):0.5\n",
      "('##i', '##k'):0.5\n",
      "('##k', '##e'):0.125\n",
      "('t', '##o'):0.3333333333333333\n",
      "('e', '##a'):0.3333333333333333\n",
      "('##a', '##t'):0.16666666666666666\n",
      "('a', '##p'):0.125\n"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T06:51:17.785737Z",
     "start_time": "2025-04-01T06:51:17.781910Z"
    }
   },
   "cell_type": "code",
   "source": [
    "best_pair = \"\"\n",
    "max_score = None\n",
    "for pair, score in pair_scores.items():\n",
    "    if max_score is None or max_score < score:\n",
    "        best_pair = pair\n",
    "        max_score = score\n",
    "print(best_pair, max_score)"
   ],
   "id": "147bdc929f1bfef0",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('S', '##h') 1.0\n"
     ]
    }
   ],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T07:00:33.474409Z",
     "start_time": "2025-04-01T07:00:33.470586Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def merge_pair(a, b, splits):\n",
    "    \"\"\"\n",
    "    在所有单词的拆分中合并指定的字符对(a, b)\n",
    "\n",
    "    参数:\n",
    "        a (str): 要合并的第一个字符/子词\n",
    "        b (str): 要合并的第二个字符/子词\n",
    "        splits (dict): 当前所有单词的拆分字典 {单词: 拆分后的字符/子词列表}\n",
    "\n",
    "    返回:\n",
    "        更新后的splits字典，其中所有(a, b)对已被合并\n",
    "    \"\"\"\n",
    "    for word in stats:\n",
    "        split = splits[word]\n",
    "        if len(split) == 1:\n",
    "            continue\n",
    "        i = 0\n",
    "        while i < len(split) - 1:\n",
    "            if split[i] == a and split[i + 1] == b:     # 找到匹配的字符对\n",
    "                merge = a + b[2:] if b.startswith(\"##\") else a + b\n",
    "                split = split[:i] + [merge] + split[i + 2:]\n",
    "            else:\n",
    "                i += 1\n",
    "        splits[word] = split    # 更新split\n",
    "    return splits"
   ],
   "id": "9e7e174a0de6905",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T07:14:21.312883Z",
     "start_time": "2025-04-01T07:14:21.307588Z"
    }
   },
   "cell_type": "code",
   "source": [
    "vocab_size = 50\n",
    "while len(vocab) < vocab_size:\n",
    "    # 1. 计算所有可能字符对的得分\n",
    "    scores = compute_pair_scores(splits)\n",
    "\n",
    "    # 2. 找出得分最高的字符对\n",
    "    best_pair, max_score = \"\", None\n",
    "    for pair, score in scores.items():\n",
    "        if max_score is None or max_score < score:\n",
    "            best_pair = pair\n",
    "            max_score = score\n",
    "\n",
    "    # 3. 在所有单词拆分中合并这个最佳字符对\n",
    "    splits = merge_pair(*best_pair, splits)\n",
    "    # 4. 创建新token并加入词表\n",
    "    new_token = (\n",
    "        best_pair[0] + best_pair[1][2:]     # 处理带##的后缀情况\n",
    "        if best_pair[1].startswith(\"##\")\n",
    "        else best_pair[0] + best_pair[1]\n",
    "    )\n",
    "    vocab.append(new_token)\n",
    "print(\"vocab:\", vocab)  # 最终的词表"
   ],
   "id": "5756623a9e4bc84d",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab: ['##a', '##e', '##g', '##h', '##i', '##k', '##l', '##o', '##p', '##r', '##s', '##t', '##u', '##v', '##y', '##果', '##欢', '##派', 'I', 'S', 'a', 'c', 'e', 'g', 'h', 'l', 't', 'v', 'y', '不', '他', '吃', '喜', '我', '苹', 'Sh', '喜欢', '苹果', '苹果派', 'li', 'lik', 'gi', 'giv', '##pl', '##ppl', '##ry', 'to', 'yo', 'ea', 'eat']\n"
     ]
    }
   ],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-01T09:20:47.033548Z",
     "start_time": "2025-04-01T09:20:47.026366Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "9209625ef90f927f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "b29fa117b04253ed"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
